#!/bin/bash

# Automatically create, rotate, and destroy periodic Proxmox VE snapshots.
# Author: Daniele Corsini <daniele.corsini@enterpriseve.com>

declare -r VERSION=0.1.8
declare -r NAME=$(basename "$0")
declare -r PROGNAME=${NAME%.*}

declare -r QEMU_CMD='qm'
declare -r LXC_CMD='pct'

declare opt_vm_ids=''
declare opt_exclude_vm_ids=''
declare -i opt_vm_state=0
declare -i opt_keep=1
declare -i opt_syslog=0
declare -i opt_debug=0
declare -i opt_dry_run=0
declare opt_label=''
declare opt_script=''

declare snap_name_prefix=''
declare -i vm_id=0
declare vm_tecnology=''
declare vm_ids
declare -r cron_file="/etc/pve/eve/autosnap/$PROGNAME.cron"  

function usage(){
    shift

    if [ "$1" != "--no-logo" ]; then
        cat << EOF
    ______      __                       _              _    ________
   / ____/___  / /____  _________  _____(_)_______     | |  / / ____/
  / __/ / __ \/ __/ _ \/ ___/ __ \/ ___/ / ___/ _ \    | | / / __/
 / /___/ / / / /_/  __/ /  / /_/ / /  / (__  )  __/    | |/ / /___
/_____/_/ /_/\__/\___/_/  / .___/_/  /_/____/\___/     |___/_____/
                         /_/
                         
EOF
    fi

    cat << EOF
EnterpriseVE automatic snapshot for Proxmox VE      (Made in Italy)

Usage:
    $PROGNAME <COMMAND> [ARGS] [OPTIONS]
    $PROGNAME help
    $PROGNAME version

    $PROGNAME create  --vmid=<string> --label=<string> --keep=<integer>
                             --vmstate --script=<string> --syslog
                             --exclude-vmid=<string>
    $PROGNAME destroy --vmid=<string> --label=<string>
    $PROGNAME enable  --vmid=<string> --label=<string>
    $PROGNAME disable --vmid=<string> --label=<string>

    $PROGNAME status
    $PROGNAME clean   --vmid=<string> --label=<string> --keep=<integer>

    $PROGNAME snap    --vmid=<string> --label=<string> --keep=<integer>
                             --vmstate --script=<string> --syslog
                             --exclude-vmid=<string>

Commands:
    version                  Show version program.
    help                     Show help program.
    create                   Create snap job from scheduler.
    destroy                  Remove snap job from scheduler.
    enable                   Enable snap job from scheduler.
    disable                  Disable snap job from scheduler.
    status                   Get list of all auto snapshots.
    clean                    Remove all auto snapshots.
    snap                     Will snap one time.

Options:
    --vmid=string            The ID of the VM/CT, comma separated (es. 100,101,102), 
                             'all' for all VM/CT in node.
    --exclude-vmid           Exclude ID of the VM/CT, comma separated (es. 100,101,102).
    --vmstate                Save the vmstate only qemu.
    --label=string           Is usually 'hourly', 'daily', 'weekly', or 'monthly'.
    --keep=integer           Specify the number of snapshots which should will keep. Default 1.
    --script=string          Use specified hook script.
                             Es. /usr/share/doc/$PROGNAME/examples/script-hook.sh
    --syslog                 Write messages into the system log.

Report bugs to <support@enterpriseve.com>.
EOF

    exit 1
}

function log(){
    local level=$1
    shift 1
    local message=$*

    case $level in
        debug) [ $opt_debug -eq 1 ] && echo -e "$(date "+%F %T") DEBUG: $message";;
        
        info) 
            echo -e "$message"; 
            [ $opt_syslog -eq 1 ] && logger -t "$PROGNAME" "$message"
            ;;
        
        error)
            echo "ERROR: $message" 1>&2
            [ $opt_syslog -eq 1 ] && logger -t "$PROGNAME" -p daemon.err "$message"
            ;;

        *)  
            echo "$message" 1>&2
            [ $opt_syslog -eq 1 ] && logger -t "$PROGNAME" "$message"
            ;;
    esac
}

function parse_opts(){
    shift

    local args
    args=$(getopt \
           --options '' \
           --longoptions=vmstate,all,vmid:,label:,keep:,exclude-vmid: \
           --longoptions=script:,syslog,debug,dry-run \
           --name "$PROGNAME" \
           -- "$@") \
           || exit 128

    eval set -- "$args"

    while true; do    
      case "$1" in
        --vmid) opt_vm_ids=$2; shift 2;;
        --exclude-vmid) opt_exclude_vm_ids=$2; shift 2;;        
        --vmstate) opt_vm_state=1; shift;;
        --label) opt_label="$2"; shift 2;;
        --keep) opt_keep=$2; shift 2;;
        --script) opt_script="$2"; shift 2;;
        --syslog) opt_syslog=1; shift;;
        --debug) opt_debug=1; shift;;
        --dry-run) opt_dry_run=1; shift;;
        --) shift; break;;
        *) break;;
      esac
    done

    [ -z "$opt_vm_ids" ] && { log info "VM id is not set."; exit 1; }
    if [ "$opt_vm_ids" = "all" ]; then
        #all VM/CT
        vm_ids=$(echo "$(qm list | awk '{print $1}' | sed 1d) $(pct list | awk '{print $1}' | sed 1d)" | tr ' ' '\n')
    else
        #comma separated
        vm_ids=$(echo "$opt_vm_ids" | tr ',' '\n')
    fi
    
    #exclude id comma separated
    local exclude_vm_ids; exclude_vm_ids=$(echo "$opt_exclude_vm_ids" | tr ',' '\n')

    #VM/CT local host
    local all_vms
    all_vms=$(echo "$(qm list | awk '{print $1}' | sed 1d) $(pct list | awk '{print $1}' | sed 1d)" | tr ' ' '\n')

    local vm_ids_tmp=""
    for vm_id in $vm_ids; do
        local vm_id_tmp
        local skip_id=0

        #exclude id if exists in array
        for vm_id_tmp in $exclude_vm_ids; do
            if [ "$vm_id_tmp" == "$vm_id" ]; then
                skip_id=1
                break
            fi
        done

        [ $skip_id -eq 1 ] && continue;

        #check valid id
        for vm_id_tmp in $all_vms; do
            if [ "$vm_id_tmp" == "$vm_id" ]; then
                [ ! -z "$vm_ids_tmp" ] && vm_ids_tmp="$vm_ids_tmp,"
                vm_ids_tmp="$vm_ids_tmp$vm_id"
                break
            fi
        done
    done
    vm_ids=$(echo "$vm_ids_tmp" | tr ',' '\n')

    [ -z "$opt_label" ] && { log info "Label is not set correctly."; exit 1; }
    [ "$opt_keep" -le 0 ] && { log info "Keep is not set correctly. Value > 0."; exit 1; }

    snap_name_prefix="auto$opt_label"
}

function get_tecnology(){
    if $QEMU_CMD list | awk '{print $1}' | grep -q $vm_id; then 
        vm_tecnology=$QEMU_CMD
    elif $LXC_CMD list | awk '{print $1}' | grep -q $vm_id; then
        vm_tecnology=$LXC_CMD
    else
        log error "VM $vm_id - Unknown tecnology"
        vm_tecnology=''
    fi
}

function fix_cron(){
     #create directory
    mkdir -p /etc/pve/eve/autosnap

    local -r cron_d="/etc/cron.d/$PROGNAME"
  
    if [[ ! -h "$cron_d" && ! -e "$cron_file" ]]; then
        #create cron file if not exist
        cat > "$cron_file" << EOL
#Cron file for $PROGNAME automatically generated
SHELL=/bin/sh
PATH=/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin

EOL
    
    elif [[ ! -h "$cron_d" ]]; then
        #not a symbolic link
        if [[ -e "$cron_file" ]]; then
            #exist configuration
            #archive old configuration
            mv "$cron_d" "$cron_file.old.$(hostname)"
        else
            mv "$cron_d" "$cron_file"
        fi        
    fi

    #symplic link
    ln -s "$cron_file" "$cron_d" 2>/dev/null
}

function cron_action_job(){
    local action=$1
    parse_opts "$@"

    local -r job_key_cron="snap --vmid=$opt_vm_ids --label='$opt_label'"
    local -i check=0; [ "$action" = "create" ] && check=1 
    if grep -h "$job_key_cron" "$cron_file"; then
        [ $check -eq 1 ] && { log info "Job already exists in cron file '$cron_file'"; exit 1; }
    else
        [ $check -eq 0 ] && { log info "Job not exists in cron file '$cron_file'"; exit 1; }
    fi

    #action
    case $action in
        create)
            local cron_scheduling="0 0 * * *"             #Default run once a day at midnight

            case "$opt_label" in
                (*hourly*) cron_scheduling="0 * * * *";;  #Run once an hour at the beginning of the hour	
                (*daily*) cron_scheduling="0 0 * * *";;   #Run once a day at midnight             
                (*weekly*) cron_scheduling="0 0 * * 0";;  #Run once a week at midnight on Sunday morning	
                (*monthly*) cron_scheduling="0 0 1 * *";; #Run once a month at midnight of the first day of the month
            esac

            local job="$cron_scheduling root $PROGNAME $job_key_cron --keep=$opt_keep"

            [ $opt_vm_state -eq 1 ] && job="$job --vmstate"                 #state memory    
            [ -e "$opt_script" ] && job="$job --script='$opt_script'"       #script hook
            [ $opt_syslog -eq 1 ] && job="$job --syslog"                    #syslog
        
            echo -e "$job" >> "$cron_file"
            ;;

        destroy) sed -i "\?$job_key_cron?d" "$cron_file";;
        enable) sed -i "\?$job_key_cron?s?^#??g" "$cron_file";;
        disable) sed -i "\?$job_key_cron?s?^?#?g" "$cron_file";;
    esac

    echo -e "Job $action in cron file '$cron_file'";
}

function clean(){
    parse_opts "$@"

    local begin=$SECONDS
    log debug "$PROGNAME $VERSION"
    log debug "Command line: $*"

    call_hook_script "clean-job-start" "-"
    
    for vm_id in $vm_ids; do 
        get_tecnology
        [ -z "$vm_tecnology" ] && continue;

        remove_old_snapshots
    done

    vm_tecnology=''
    vm_id=0

    call_hook_script "clean-job-end" "-"

    log debug "Execution: $((SECONDS-begin)) sec."
}

function remove_old_snapshots(){
    [ -z "$vm_tecnology" ] && return;

    local snapshots;
    snapshots=$($vm_tecnology listsnapshot $vm_id | \
                awk '{print $1}' | \
                grep "$snap_name_prefix" | sort -r)

    local -i index=0
    local snap_name=''
    for snap_name in $snapshots; do
        if [ "$index" -ge "$opt_keep" ]; then            
            log info "VM $vm_id - Removing snapshot $snap_name"

            call_hook_script "snap-remove-pre" "$snap_name"

            if ! do_run "$vm_tecnology delsnapshot $vm_id $snap_name"; then
                call_hook_script "snap-remove-abort" "$snap_name"
                log error "VM $vm_id - Abort remove snapshot $snap_name"
                continue;
            fi

            call_hook_script "snap-remove-post" "$snap_name"
        fi
        
        let index++
    done
}

function status(){
    local print_header=1

    local tecnology=''
    for tecnology in $QEMU_CMD $LXC_CMD; do
        local vm=''
        for vm in $($tecnology list | awk '{print $1}' | sed 1d); do
            local result;
            result=$($tecnology listsnapshot "$vm" | \
                     grep "$PROGNAME" | \
                     sort -r | \
                     awk -v vm="$vm" '{print vm,substr($1,1+length($1)-12),substr($1,5,1+length($1)-12-5)}' | \
                     awk 'BEGIN {FIELDWIDTHS="4 2 2 2 2 2 2 10"} {printf "%s %s-%s-%s %s:%s:%s %s\n",$1,$2,$3,$4,$5,$6,$7,$8}')

            if [ ! -z "$result" ]; then
                if [ $print_header -eq 1 ]; then
                    print_header=0
                    echo "VM   SNAPSHOTS          LABEL"
                fi

                echo "${result}"
            fi
        done
    done
}

function call_hook_script(){
    export EVE4PVE_AUTOSNAP_PHASE="$1"
    export EVE4PVE_AUTOSNAP_SNAP_NAME="$2"
    export EVE4PVE_AUTOSNAP_VMID=$vm_id
    export EVE4PVE_AUTOSNAP_VMTECNOLOGY=$vm_tecnology
    export EVE4PVE_AUTOSNAP_LABEL="$opt_label"
    export EVE4PVE_AUTOSNAP_KEEP=$opt_keep
    export EVE4PVE_AUTOSNAP_VMSTATE=$opt_vm_state

    log debug "------------------------------------------------------------" 
    log debug "EVE4PVE_AUTOSNAP_PHASE:        $EVE4PVE_AUTOSNAP_PHASE" 
    log debug "EVE4PVE_AUTOSNAP_VMID:         $EVE4PVE_AUTOSNAP_VMID" 
    log debug "EVE4PVE_AUTOSNAP_VMTECNOLOGY:  $EVE4PVE_AUTOSNAP_VMTECNOLOGY" 
    log debug "EVE4PVE_AUTOSNAP_LABEL:        $EVE4PVE_AUTOSNAP_LABEL" 
    log debug "EVE4PVE_AUTOSNAP_KEEP:         $EVE4PVE_AUTOSNAP_KEEP" 
    log debug "EVE4PVE_AUTOSNAP_VMSTATE:      $EVE4PVE_AUTOSNAP_VMSTATE" 
    log debug "EVE4PVE_AUTOSNAP_SNAP_NAME:    $EVE4PVE_AUTOSNAP_SNAP_NAME" 
    log debug "------------------------------------------------------------" 

    if [ -e "$opt_script" ]; then
        log debug "VM $vm_id - Script hook: $opt_script" 
        do_run "$opt_script"
    fi
}

function snap(){
    parse_opts "$@"

    log debug "$PROGNAME $VERSION"
    log debug "Command line: $*"

    local begin=$SECONDS
    local snap_name; snap_name="$snap_name_prefix$(date +%y%m%d%H%M%S)"

    call_hook_script "snap-job-start" "-"

    for vm_id in $vm_ids; do
        get_tecnology        
        [ -z "$vm_tecnology" ] && continue;

        local vm_state='';
        if [ $vm_tecnology = $QEMU_CMD ]; then
            #memory status
            [ $opt_vm_state -eq 1 ] && vm_state="--vmstate"

            #agent enabled optimize freeze the guest file system
            if ! $vm_tecnology config "$vm_id" --current | grep 'agent: 1' 1>/dev/null; then
                log info "VM $vm_id consider enabling QEMU agent see https://pve.proxmox.com/wiki/Qemu-guest-agent"        
            fi
        fi

        #create snapshot
        call_hook_script "snap-create-pre" "$snap_name"

        log info "VM $vm_id - Creating snapshot $snap_name"
        if ! do_run "$vm_tecnology snapshot $vm_id $snap_name -description '$PROGNAME' $vm_state"; then
            call_hook_script "snap-create-abort" "$snap_name"
            log error "VM $vm_id - Create snapshot $snap_name"
            continue;
        fi

        call_hook_script "snap-create-post" "$snap_name"

        #remove old snapshot
        remove_old_snapshots    
    done

    vm_tecnology=''
    vm_id=0

    call_hook_script "snap-job-end" "-"

    log debug "Execution: $((SECONDS-begin)) sec."
}

function do_run(){
    local cmd=$*;
    local -i rc=0;

    if [ $opt_dry_run -eq 1 ]; then
        echo "$cmd"
        rc=$?
    else
        log debug "$cmd"
        eval "$cmd"
        rc=$?
        [ $rc != 0 ] && log error "$cmd"
        log debug "return $rc"
    fi

    return $rc
}

function main(){    
    [ $# = 0 ] && usage;

    fix_cron

    case "$1" in
        version) echo "$VERSION";;
        help) usage "$@";;
        create|destroy|enable|disable) cron_action_job "$@";;
        clean) clean "$@";;
        status) status;; 
        snap) snap "$@";;
        *) usage;;
    esac

    exit 0;
}

main "$@"